from typing import List, Dict, Any, Optional
from python_depend.gms_api_client import get_quote_kline
import numpy as np
import pandas as pd
import pymysql

from python_depend.mysql_pool import MySQLPool


# ==============================
# 数据库操作函数
# ==============================
def fetch_stock_list() -> List[Dict[str, Any]]:
    mysql_pool = MySQLPool.get_pool('ainvestdb')
    """从数据库获取股票列表"""
    sql_query = """
    SELECT
        CAST(thsmarket_code_hq + 256 AS SIGNED) AS market,
        thscode_hq AS code
    FROM
        pub205
    WHERE
        F005V_PUB205 IN ('212049', '212050', '212010', '212011', '212210')
        AND f018v_pub205 = '正常上市'
        AND F007D_PUB205 IS NOT NULL
        AND thscode_hq IS NOT NULL
        AND F003V_PUB205 IN ('海外普通股', '优先股', '存托凭证(普通股)', '存托凭证(优先股)', 'UNIT');
    """
    try:
        with mysql_pool.connection() as connection:
            with connection.cursor() as cursor:
                cursor.execute(sql_query)
                return cursor.fetchall()
    except pymysql.Error as e:
        print(f"[ERROR] 获取股票列表失败: {e}")
        return []


# ==============================
# 数据写入函数
# ==============================
def write_to_mysql(indicator, df, market, symbol):
    """将失衡区域信息保存到数据库"""
    imbalances = indicator.get_imbalances()
    if not imbalances:
        print("没有失衡区域需要保存")
        return
    values = []
    for imbalance in imbalances:
        trade_date = imbalance.date
        imbalances_sign = "bearish"
        if imbalance.bullish:
            imbalances_sign = "bullish"
        gap_flag = 1 if imbalance.is_gap else 0

        values.append((
            market,
            symbol,
            trade_date,
            imbalances_sign,
            str(imbalance.top),
            str(imbalance.bottom),
            str(df.loc[trade_date].open),
            str(df.loc[trade_date].high),
            str(df.loc[trade_date].low),
            str(df.loc[trade_date].close),
            str(df.loc[trade_date].volume),
            gap_flag
        ))
    mysql_pool = MySQLPool.get_pool('ext_quote_data')
    insert_query = """
                      INSERT INTO indicator_volume_gaps_imbalances 
                      (market, symbol, trade_date, imbalances_sign, top, bottom, 
                       open, high, low, close, volume, is_gap, create_time, update_time)
                      VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, now(), now())
                      ON DUPLICATE KEY UPDATE
                      top = VALUES(top), 
                      bottom = VALUES(bottom), 
                      is_gap = VALUES(is_gap), 
                      update_time = VALUES(update_time)
                      """
    try:
        connection = mysql_pool.connection()
        with connection.cursor() as cursor:
            cursor.executemany(insert_query, values)
            connection.commit()
            print(f"成功写入 {len(values)} 条数据")
    except Exception as e:
        print(f"写入失败: {e}")


# ==============================
# 数据写入函数
# ==============================
def delete_data(market, symbol):
    mysql_pool = MySQLPool.get_pool('ext_quote_data')
    delete_sql = f"delete from indicator_volume_gaps_imbalances where symbol = '{symbol}' and market = '{market}'"
    try:
        connection = mysql_pool.connection()
        with connection.cursor() as cursor:
            cursor.execute(delete_sql)
            connection.commit()
            print(f"删除{symbol}历史数据成功")
    except Exception as e:
        print(f"删除失败失败: {e}")


def convert_to_dataframe(data: Dict[str, Any]) -> pd.DataFrame:
    """将API返回数据解析为DataFrame"""
    if data.get('data') is None:
        return pd.DataFrame()
    stock_data = data.get('data', {}).get('quote_data', [{}])[0]
    data_fields = stock_data.get('data_fields', [])
    value_list = stock_data.get('value', [])
    # 字段映射与校验
    field_mapping = {
        'trade_day': '65558',  # 交易日
        'open': '7',
        'high': '8',
        'low': '9',
        'close': '11',
        'volume': '13'
    }

    missing_fields = [f for f in field_mapping.values() if f not in data_fields]
    if missing_fields:
        print(f"[WARNING] 缺失关键字段: {missing_fields}，跳过数据转换")
        return pd.DataFrame()

    # 提取字段索引
    idx_map = {key: data_fields.index(value) for key, value in field_mapping.items()}

    # 构建DataFrame
    try:
        df = pd.DataFrame({
            'open': [float(values[idx_map['open']]) if values[idx_map['open']] else np.nan for values in value_list],
            'high': [float(values[idx_map['high']]) if values[idx_map['high']] else np.nan for values in value_list],
            'low': [float(values[idx_map['low']]) if values[idx_map['low']] else np.nan for values in value_list],
            'close': [float(values[idx_map['close']]) if values[idx_map['close']] else np.nan for values in value_list],
            'volume': [float(values[idx_map['volume']]) if values[idx_map['volume']] else np.nan for values in
                       value_list],
        }, index=[
            values[idx_map['trade_day']] for values in value_list
        ])
        df.index.name = 'trade_day'
        return df.dropna()  # 过滤无效数据
    except Exception as e:
        print(f"[ERROR] 数据转换失败: {e}")
        return pd.DataFrame()


# ==============================
# K线数据处理函数
# ==============================
def fetch_kline_data(market: str, code: str, count: int = 200) -> Optional[pd.DataFrame]:
    """获取K线数据并转换为DataFrame"""

    request_params = {
        "code_list": [{"market": market, "codes": [code]}],
        "trade_class": "intraday",
        "time_period": "day_1",
        "time_range": {"count": count, "end_time": 0},
        "adjust_type": "forward"
    }
    try:
        result = get_quote_kline(request_params)
        return convert_to_dataframe(result)
    except Exception as e:
        print(f"[ERROR] 获取{market}/{code} K线数据失败: {e}")
        return pd.DataFrame()


class Imbalance:
    def __init__(self, bullish, date, top, bottom, is_gap=False):
        self.bullish = bullish  # 是否为看涨失衡
        self.date = date  # 失衡发生的日期 (datetime对象)
        self.top = top  # 失衡区域的顶部价格
        self.bottom = bottom  # 失衡区域的底部价格
        self.is_gap = is_gap  # 是否为价格缺口


# ==============================
# 核心算法逻辑
# ==============================
class VolumeGapsAndImbalances:
    def __init__(self, show_limit=10):
        # 初始化参数
        self.vi_show_limit = show_limit  # 显示的最大失衡数量
        self.imbe = []  # 看跌失衡列表
        self.imbu = []  # 看涨失衡列表

    def cleanup(self, imbalance_list, bullish, current_high, current_low):
        """清理不再有效的失衡区域"""
        while imbalance_list and ((bullish and current_low <= imbalance_list[-1].top) or
                                  (not bullish and current_high >= imbalance_list[-1].top)):
            imbalance_list.pop()

    def detect_imbalances(self, df):
        """检测DataFrame中的失衡区域和缺口"""
        # 确保数据有足够的长度
        if len(df) < 2:
            return

        # 遍历数据，从第二根K线开始
        for i in range(1, len(df)):
            # 获取当前和前一根K线的数据
            current_bar = df.iloc[i]
            prev_bar = df.iloc[i - 1]

            # 获取日期
            prev_date = df.index[i - 1]  # 前一根K线的日期

            # 计算当前和前一根K线的最高价和最低价
            current_high = max(current_bar['open'], current_bar['close'])
            current_low = min(current_bar['open'], current_bar['close'])
            prev_high = max(prev_bar['open'], prev_bar['close'])
            prev_low = min(prev_bar['open'], prev_bar['close'])

            # 清理不再有效的失衡区域
            self.cleanup(self.imbe, False, current_high, current_low)
            self.cleanup(self.imbu, True, current_high, current_low)

            # 检测看跌失衡
            if prev_low > current_high:
                is_gap = prev_bar['low'] - current_bar['high'] > 0
                # 使用前一根K线的日期，与原代码逻辑一致
                self.imbe.append(Imbalance(False, prev_date, prev_low, current_high, is_gap))

            # 检测看涨失衡
            elif prev_high < current_low:
                is_gap = current_bar['low'] - prev_bar['high'] > 0
                # 使用前一根K线的日期，与原代码逻辑一致
                self.imbu.append(Imbalance(True, prev_date, prev_high, current_low, is_gap))

    def get_imbalances(self):
        """获取检测到的失衡区域"""
        return self.imbu + self.imbe


def execute_job(k_count):
    """主任务执行流程"""
    print("[START] 开始执行任务")
    # 1. 获取股票列表
    stock_list = fetch_stock_list()
    # stock_list = [{"code": "AAPL", "market": 185}]
    if not stock_list:
        print("[ERROR] 股票列表为空，任务终止")
        return
    print("查询股票数量：" + str(len(stock_list)))
    # 2. 处理每只股票
    for stock in stock_list:
        market = str(stock['market'])
        code = stock['code']
        print(f"\n[PROCESS] 处理股票: Market={market}, Code={code}")
        # 获取K线数据
        df = fetch_kline_data(market=market, code=code, count=k_count)
        if df.empty:
            print(f"[WARNING] {market}/{code} 无有效K线数据")
            continue
        indicator = VolumeGapsAndImbalances()
        indicator.detect_imbalances(df)

        # 获取并打印检测结果
        result = indicator.get_imbalances()

        for imbalance in result:
            print(
                f"看涨失衡：{imbalance.bullish} 日期: {imbalance.date}, 区域: {imbalance.top} - {imbalance.bottom}, 是否为缺口: {imbalance.is_gap}")
        print("结果写入数据库:")
        #删除历史计算数据
        delete_data(market,code)
        write_to_mysql(indicator, df, market, code)
    print("[END] 任务执行完成")


# 使用示例
if __name__ == "__main__":
    MySQLPool.initialize()
    print("开始执行任务")
    execute_job(200)
    print("关闭数据库连接")
    MySQLPool.close_all()
